#!/usr/bin/python
# -*- encoding: utf-8 -*-
from libs.loss import OhemCELoss
from libs.optimizer import Optimizer
from libs.utils import save_sample_imgs, showimg, showmask, batch_end_log, epoch_end_log
import torch
import os
from libs.loss import iou_loss, HairMattingLoss, bineary_iou_loss
from trainer.BaseTrainer import BaseTrainer
from tqdm import tqdm

class Trainer(BaseTrainer):
    def __init__(self, userconfig, train_data_loader, test_data_loader, checkpoint_dir, sample_dir):
        super(Trainer, self).__init__(userconfig, train_data_loader, test_data_loader, checkpoint_dir, sample_dir)

    def eval(self):
        print("star to test...")
        self.net.eval()
        with torch.no_grad():
            self.reset_test_monitor()

            for samples in tqdm(self.test_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask_gt'].to(self.device)
                mask_gt = mask_gt.long()
                mask_gt = torch.squeeze(mask_gt, 1)

                mask_pred, mask_pred_16, mask_pred_32 = self.net(image)

                loss1 = self.Loss1(mask_pred, mask_gt)
                loss2 = self.Loss2(mask_pred_16, mask_gt)
                loss3 = self.Loss3(mask_pred_32, mask_gt)

                test_loss = loss1 + loss2 + loss3
                test_loss_value = test_loss.item()
                self.test_loss_monitor.update(test_loss_value)

                mask_pred = torch.nn.functional.softmax(mask_pred)
                mask_pred = mask_pred[:, 1, :, :]
                mask_gt = mask_gt.float()

                aiou = iou_loss(mask_pred, mask_gt)
                self.test_aiou_monitor.update(aiou.item())

            avg_loss = self.test_loss_monitor.avg
            avg_aiou = self.test_aiou_monitor.avg

            print("test phase | loss: %.4f aiou: %.4f" % (avg_loss, avg_aiou))

        return avg_loss, avg_aiou


    def train(self):
        ## optimizer
        momentum = 0.9
        weight_decay = 5e-4
        lr_start = 1e-2
        max_iter = 80000
        power = 0.9
        warmup_steps = 1000
        warmup_start_lr = 1e-5
        optim = Optimizer(
                model=self.net,
                lr0=lr_start,
                momentum=momentum,
                wd=weight_decay,
                warmup_steps=warmup_steps,
                warmup_start_lr=warmup_start_lr,
                max_iter=max_iter,
                power=power)

        # model
        ignore_idx = -100
        score_thres = 0.7
        n_min = self.userconfig.batch_size * self.userconfig.input_shape[0] * self.userconfig.input_shape[1]//16

        self.Loss1 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
        self.Loss2 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)
        self.Loss3 = OhemCELoss(thresh=score_thres, n_min=n_min, ignore_lb=ignore_idx)

        for epoch in range(self.userconfig.max_epoch):
            self.reset_train_monitor()
            self.net.train()

            for step, samples in enumerate(self.train_data_loader):
                image = samples['image'].to(self.device)
                mask_gt = samples['mask_gt'].to(self.device)

                mask_gt = mask_gt.long()
                mask_gt = torch.squeeze(mask_gt, 1)
                optim.zero_grad()
                mask_pred, mask_pred_16, mask_pred_32 = self.net(image)

                loss1 = self.Loss1(mask_pred, mask_gt)
                loss2 = self.Loss2(mask_pred_16, mask_gt)
                loss3 = self.Loss3(mask_pred_32, mask_gt)
                train_loss = loss1 + loss2 + loss3
                train_loss_value = train_loss.item()
                self.train_loss_monitor.update(train_loss_value)

                mask_pred = torch.nn.functional.softmax(mask_pred)
                mask_pred = mask_pred[:, 1, :, :]

                mask_gt = mask_gt.float()

                aiou_value = iou_loss(mask_pred, mask_gt).item()
                self.train_aiou_monitor.update(aiou_value)

                train_loss.backward()
                optim.step()

                batch_end_log(log_intervel=self.userconfig.log_intervel, max_epoch=self.userconfig.max_epoch,
                              epoch=epoch, step=step, max_step=self.train_image_len,lr=optim.lr,
                              train_loss=train_loss_value, aiou=aiou_value)

                if step % self.userconfig.sample_intervel == 0:
                    save_sample_imgs(samples, {"mask_pred": mask_pred, "dyemask": mask_pred}, self.sample_dir, epoch, step)


            epoch_end_log(phase="train", train_loss=self.train_loss_monitor.avg, aiou=self.train_aiou_monitor.avg)
            self.eval()
            torch.save(self.net.state_dict(), os.path.join(self.checkpoint_dir, self.ckpt_name()))
